import numpy as np
import matplotlib.pyplot as plt

# 设置中文字体（可选，解决中文显示问题）和整体风格
plt.rcParams['font.sans-serif'] = ['SimHei'] # 如果用中文标签
plt.rcParams['axes.unicode_minus'] = False   # 正常显示负号
plt.style.use('default') # 使用默认白色背景风格

# 1. 创建数据点
# 避开 p=0 和 p=1 这两个导致log(0)的点
p = np.linspace(0.001, 0.999, 500) # 在0.001到0.999之间生成500个点

# 2. 计算熵值 H(p) = -p*log2(p) - (1-p)*log2(1-p)
# 使用 np.log2，并以 np.clip 避免数值计算错误（虽是冗余安全措施）
H = -p * np.log2(np.clip(p, 1e-10, 1)) - (1-p) * np.log2(np.clip(1-p, 1e-10, 1))

# 3. 创建图形和坐标轴
fig, ax = plt.subplots(figsize=(8, 6)) # 设置图片大小

# 4. 绘制曲线
ax.plot(p, H, 'k-', linewidth=2) # 'k-' 表示黑色实线，linewidth设置线宽

# 5. 设置坐标轴标签和标题（完全按照您的图片）
ax.set_xlabel('$P_{\\oplus}$', fontsize=12) # 使用 LaTeX 语法表示下标 P⊕
ax.set_ylabel('Entropy(S)', fontsize=12)
# ax.set_title('Entropy Function') # 如果需要标题可以取消注释

# 6. 设置坐标轴范围
ax.set_xlim(0.0, 1.0)
ax.set_ylim(0.0, 1.0)

# 7. 设置刻度，与参考图保持一致
ax.set_xticks([0.0, 0.5, 1.0])
ax.set_yticks([0.0, 0.5, 1.0])
# 可以设置更密集的刻度网格
# ax.grid(True, linestyle='--', alpha=0.7) # 如果需要网格线可以取消注释

# 8. 在曲线末端添加箭头（可选，模仿原图）
# 在 p=0.99 附近取点计算方向向量，并添加箭头
idx = 400 # 在曲线后半部分选取一个点
ax.annotate('', xytext=(p[idx], H[idx]), xy=(p[idx+10], H[idx+10]),
            arrowprops=dict(arrowstyle='->', color='black', lw=1))

# 9. 美化图形（使图形紧凑）
plt.tight_layout()

# 10. 显示图形
plt.show()

# 11. 保存图形（取消注释以下行以保存为图片文件）
# fig.savefig('entropy_curve.png', dpi=300, bbox_inches='tight') # 保存为高清PNG
# fig.savefig('entropy_curve.pdf', bbox_inches='tight') # 保存为矢量PDF，用于论文